// ================================================
// CodeGenerator.ts
// ================================================

import {
  PipelineService, Node, Flow
} from './PipelineService';
import { BaseCodeGenerator, NodeObject } from './BaseCodeGenerator';

export class CodeGenerator extends BaseCodeGenerator {

  static generateCodeForNodes(
    flow: Flow,
    componentService: any,
    targetNodeId: string,
    fromStart: boolean,
    variablesAutoNaming: boolean
  ): {
    codeList: string[];
    incrementalCodeList: { code: string; nodeId: string }[];
    executedNodes: Set<string>;
  } {
    // Similar logic as your original "generateCodeForNodes" method:
    const codeList: string[] = [];
    const incrementalCodeList: { code: string; nodeId: string }[] = [];
    const executedNodes = new Set<string>();
    const uniqueImports = new Set<string>();
    const uniqueDependencies = new Set<string>();
    const functions = new Set<string>();
    const envMap = new Map<string, Node>();
    const connMap = new Map<string, Node>();

    const { nodesToTraverse, nodesMap } = this.computeNodesToTraverse(
      flow, 
      targetNodeId, 
      componentService
    );

    // Identify special nodes
    flow.nodes.forEach(node => {
      const type = componentService.getComponent(node.type)._type;
      if (['env_variables','env_file'].includes(type)) {
        envMap.set(node.id, node);
      } else if (type === 'connection') {
        connMap.set(node.id, node);
      }
    });

    // Build node objects
    const nodeObjects = this.createNodeObjects(
      flow, 
      componentService, 
      nodesToTraverse, 
      nodesMap, 
      variablesAutoNaming
    );

    // Process
    for (const nodeObj of nodeObjects) {
      nodeObj.imports.forEach(i => uniqueImports.add(i));
      nodeObj.dependencies.forEach(d => uniqueDependencies.add(d));
      nodeObj.functions.forEach(f => functions.add(f));

      const nodeCode = [
        ...nodeObj.imports,
        ...nodeObj.functions,
        nodeObj.code
      ].join('\n');

      incrementalCodeList.push({ code: nodeCode, nodeId: nodeObj.id });
      codeList.push(nodeObj.code);
      executedNodes.add(nodeObj.id);

      // If it's the target node, handle display code or final try-catch
      if (nodeObj.id === targetNodeId) {
        let displayCode = '';
        if (nodeObj.type.includes('processor') || nodeObj.type.includes('input')) {
          if (nodeObj.type.includes('documents')) {
            if (!fromStart) {
              codeList.length = 0;
              codeList.push(nodeObj.code);
              executedNodes.clear();
              executedNodes.add(nodeObj.id);
            }
            displayCode = `\n_amphi_display_documents_as_html(${nodeObj.outputName})`;
          } else {
            if (!fromStart) {
              codeList.length = 0;
              codeList.push(nodeObj.code);
              executedNodes.clear();
              executedNodes.add(nodeObj.id);
            }
            displayCode = `\__amphi_display(${nodeObj.outputName}, dfName="${nodeObj.outputName}", nodeId="${targetNodeId}"${nodeObj.runtime !== "local" ? `, runtime="${nodeObj.runtime}"` : ''})`;
          }
          codeList.push(displayCode);
          if (incrementalCodeList.length > 0) {
            incrementalCodeList[incrementalCodeList.length - 1].code += displayCode;
          }
        } else if (nodeObj.type.includes('output')) {
          // Wrap code
          const combined = codeList.join('\n');
          const indented = combined.split('\n').map(x => `    ${x}`).join('\n');
          codeList.length = 0;
          codeList.push('try:\n' + indented);
          codeList.push('    print("Execution has been successful")\n');
          codeList.push('except Exception as e:\n');
          codeList.push('    print(f"Execution failed with error {e}")\n');
          codeList.push('    raise\n');
        }
      }
    }

    // Handle env/connection code
    let envVariablesCode = '';
    envMap.forEach(node => {
      const comp = componentService.getComponent(node.type);
      const config: any = node.data;
      envVariablesCode += comp.generateComponentCode({ config });
      comp.provideImports({ config }).forEach(i => uniqueImports.add(i));
    });

    let connectionsCode = '';
    connMap.forEach(node => {
      const comp = componentService.getComponent(node.type);
      const config: any = node.data;
      connectionsCode += comp.generateComponentCode({ config });
      comp.provideImports({ config }).forEach(i => uniqueImports.add(i));
    });

    // Final build
    const now = new Date();
    const dateString = now.toISOString().replace(/T/, ' ').replace(/\..+/, '');
    const dateComment = `# Source code generated by Amphi\n# Date: ${dateString}`;
    const additionalImports = `# Additional dependencies: ${Array.from(uniqueDependencies).join(', ')}`;

    const finalList = [
      dateComment,
      additionalImports,
      ...Array.from(uniqueImports),
      envVariablesCode,
      connectionsCode,
      ...Array.from(functions),
      ...codeList
    ].filter(Boolean);

    const formatted = finalList.map(code => this.formatVariables(code));
    return {
      codeList: formatted,
      incrementalCodeList,
      executedNodes
    };
  }

  static generateCode(
    pipelineJson: string, 
    commands: any, 
    componentService: any, 
    variablesAutoNaming: boolean
  ): string {
    const { codeList } = this.generateCodeForNodes(
      PipelineService.filterPipeline(pipelineJson),
      componentService,
      'none',
      true,
      variablesAutoNaming
    );
    return codeList.join('\n');
  }

  static generateCodeUntil(
    pipelineJson: string,
    commands: any,
    componentService: any,
    targetNode: string,
    incremental: boolean,
    variablesAutoNaming: boolean
  ): any[] {
    const flow = PipelineService.filterPipeline(pipelineJson);
    if (flow.nodes.some(n => !n.data.nameId)) {
      variablesAutoNaming = true;
    }
    const { nodesToTraverse, nodesMap } = this.computeNodesToTraverse(
      flow,
      targetNode,
      componentService
    );

    let fromStart = false;
    let maxLastUpdated = 0;

    for (const nodeId of nodesToTraverse) {
      if (nodeId === targetNode) continue;
      const node = nodesMap.get(nodeId);
      if (!node) continue;

      const data = node.data || {};
      const lastUpdated = data.lastUpdated || 0;
      const lastExecuted = data.lastExecuted || 0;

      if (lastUpdated > maxLastUpdated) {
        maxLastUpdated = lastUpdated;
      }
      if (lastUpdated >= lastExecuted) {
        fromStart = true;
        break;
      }
    }

    if (variablesAutoNaming) {
      fromStart = true;
    }

    const { codeList, incrementalCodeList, executedNodes } = this
      .generateCodeForNodes(flow, componentService, targetNode, fromStart, variablesAutoNaming);

    if (fromStart) {
      console.log("Generating code from start (fromStart: true).");
      const command = 'pipeline-metadata-panel:delete-all';
      commands.execute(command, {}).catch(reason => {
        console.error(`Error executing ${command}. Reason: ${reason}`);
      });
    } else {
      console.log("No updates in previous nodes. Generating code for target node only.");
    }

    const now = Date.now();
    executedNodes.forEach(id => {
      const node = nodesMap.get(id);
      if (node && node.data) {
        node.data.lastExecuted = now;
      }
    });
    const targetData = nodesMap.get(targetNode)?.data;
    if (targetData) {
      targetData.lastExecuted = now;
    }

    return incremental ? incrementalCodeList : codeList;
  }
}